package weekthree;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

public class LowestCommonAncestorOfBtree_236 {
    private Map<Integer, TreeNode> fatherMap;
    private Set<Integer> redNodes;

    public TreeNode lowestCommonAncestor(TreeNode root, TreeNode p, TreeNode q) {
        fatherMap = new HashMap<>();
        redNodes = new HashSet<>();
        calcFather(root);
        redNodes.add(root.val);
        while (p.val != root.val) {
            redNodes.add(p.val);
            p = fatherMap.get(p.val);
        }
        while (!redNodes.contains(q.val)) {
            q = fatherMap.get(q.val);
        }
        return q;
    }

    private void calcFather(TreeNode root) {
        if (root == null) return;

        if (root.left != null) {
            fatherMap.put(root.left.val, root);
            calcFather(root.left);
        }
        if (root.right != null) {
            fatherMap.put(root.right.val, root);
            calcFather(root.right);
        }
    }

    private class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;

        TreeNode(int x) {
            val = x;
        }
    }
}
